# Copyright Niantic 2021. Patent Pending. All rights reserved.
#
# This software is licensed under the terms of the ManyDepth licence
# which allows for non-commercial use only, the full terms of which are made
# available in the LICENSE file.

import os
import argparse


class MonoscaledepthOptions:
    def __init__(self):
        # the directory that options.py resides in
        file_dir = os.path.dirname(__file__)

        self.parser = argparse.ArgumentParser(description="MonoScaleDepth options")

        # PATHS
        self.parser.add_argument(
            "--data_path",
            type=str,
            help="path to the training data",
            default=os.path.join(os.path.expanduser("~"), "kitti_dataset"),
        )
        self.parser.add_argument(
            "--log_dir",
            type=str,
            help="log directory",
            default=os.path.join(os.path.expanduser("~"), "tmp"),
        )

        # TRAINING options
        self.parser.add_argument(
            "--model_name",
            type=str,
            help="the name of the folder to save the model in",
            default="mdp",
        )
        self.parser.add_argument(
            "--split",
            type=str,
            help="which training split to use",
            choices=[
                "kitti_raw_pose",
                "eigen_zhou",
                "odom",
                "dominant",
            ],
            default="eigen_zhou",
        )
        self.parser.add_argument(
            "--num_layers",
            type=int,
            help="number of resnet layers",
            choices=[18, 34, 50, 101, 152],
            default=18,
        )
        self.parser.add_argument(
            "--depth_binning",
            help="defines how the depth bins are constructed for the cost volume. "
            "'linear' is uniformly sampled in depth space,"
            "'inverse' is uniformly sampled in inverse depth space",
            type=str,
            choices=["linear", "inverse"],
            default="linear",
        ),
        self.parser.add_argument("--num_depth_bins", type=int, default=96)
        self.parser.add_argument(
            "--height", type=int, help="input image height", default=192
        )
        self.parser.add_argument(
            "--width", type=int, help="input image width", default=640
        )
        self.parser.add_argument(
            "--disparity_smoothness",
            type=float,
            help="disparity smoothness weight",
            default=1e-3,
        )
        self.parser.add_argument(
            "--scales",
            nargs="+",
            type=int,
            help="scales used in the loss",
            default=[0, 1, 2, 3],
        )
        self.parser.add_argument(
            "--min_depth", type=float, help="minimum depth", default=0.1
        )
        self.parser.add_argument(
            "--max_depth", type=float, help="maximum depth", default=100.0
        )
        self.parser.add_argument(
            "--frame_ids",
            nargs="+",
            type=int,
            help="frames to load",
            default=[0, -1, 1],
        )
        self.parser.add_argument(
            "--dataset",
            type=str,
            help="dataset to train on",
            choices=[
                "kitti_raw",
                "kitti_raw_pose",
                "kitti_raw_pose_semantic",
                "kitti_odom",
                "kitti_odom_pose",
                "dominant_pose",
            ],
            default="kitti_raw_pose",
        )
        self.parser.add_argument(
            "--png",
            help="if set, trains from raw KITTI png files (instead of jpgs)",
            action="store_true",
        )

        # OPTIMIZATION options
        self.parser.add_argument(
            "--batch_size", type=int, help="batch size", default=12
        )
        self.parser.add_argument(
            "--learning_rate", type=float, help="learning rate", default=1e-4
        )
        self.parser.add_argument(
            "--num_epochs", type=int, help="number of epochs", default=20
        )
        self.parser.add_argument(
            "--scheduler_step_size",
            type=int,
            help="step size of the scheduler",
            default=15,
        )
        self.parser.add_argument(
            "--freeze_teacher_and_pose",
            action="store_true",
            help="If set, freeze the weights of the single frame network and pose network.",
        )
        self.parser.add_argument(
            "--freeze_teacher_epoch",
            type=int,
            default=15,
            help="Sets the epoch number at which to freeze the teacher network and the pose network.",
        )
        self.parser.add_argument("--pytorch_random_seed", default=None, type=int)

        # ABLATION options
        self.parser.add_argument(
            "--weights_init",
            type=str,
            help="pretrained or scratch",
            choices=["pretrained", "scratch"],
            default="pretrained",
        )
        self.parser.add_argument(
            "--no_ssim", help="if set, disables ssim in the loss", action="store_true"
        )
        self.parser.add_argument(
            "--disable_automasking",
            help="if set, doesn't do auto-masking",
            action="store_true",
        )
        self.parser.add_argument(
            "--add_pose_supervise",
            action="store_true",
            help="If set, supervise pose.",
        )
        self.parser.add_argument(
            "--begin_supervise_epoch",
            type=int,
            default=0,
            help="Sets the epoch number at which to begin the pose supervise.",
        )
        self.parser.add_argument(
            "--pose_weight",
            type=float,
            help="pose weight",
            default=5e-2,
        )
        self.parser.add_argument(
            "--use_future_frame",
            action="store_true",
            help="If set, will also use a future frame in time for matching.",
        )
        self.parser.add_argument(
            "--num_matching_frames",
            help="Sets how many previous frames to load to build the cost volume",
            type=int,
            default=1,
        )
        self.parser.add_argument(
            "--disable_motion_masking",
            help="If set, will not apply consistency loss in regions where the cost volume is deemed untrustworthy",
            action="store_true",
        )
        self.parser.add_argument(
            "--no_matching_augmentation",
            action="store_true",
            help="If set, will not apply static camera augmentation or zero cost volume augmentation during training",
        )
        self.parser.add_argument(
            "--no_multi_depth",
            action="store_true",
            help="If set, only use mono depth",
        )
        self.parser.add_argument(
            "--use_semantic",
            action="store_true",
            help="If set, use semantics mask.",
        )

        # SYSTEM options
        self.parser.add_argument(
            "--no_cuda", help="if set disables CUDA", action="store_true"
        )
        self.parser.add_argument(
            "--num_workers", type=int, help="number of dataloader workers", default=12
        )

        # LOADING options
        self.parser.add_argument(
            "--load_weights_folder", type=str, help="name of model to load"
        )
        self.parser.add_argument("--mono_weights_folder", type=str)
        self.parser.add_argument(
            "--models_to_load",
            nargs="+",
            type=str,
            help="models to load",
            default=["encoder", "depth", "pose_encoder", "pose"],
        )

        # LOGGING options
        self.parser.add_argument(
            "--log_frequency",
            type=int,
            help="number of batches between each tensorboard log",
            default=250,
        )
        self.parser.add_argument(
            "--save_frequency",
            type=int,
            help="number of epochs between each save",
            default=1,
        )

        # EVALUATION options
        self.parser.add_argument(
            "--eval_stereo", help="if set evaluates in stereo mode", action="store_true"
        )
        self.parser.add_argument(
            "--eval_mono", help="if set evaluates in mono mode", action="store_true"
        )
        self.parser.add_argument(
            "--disable_median_scaling",
            help="if set disables median scaling in evaluation",
            action="store_true",
        )
        self.parser.add_argument(
            "--pred_depth_scale_factor",
            help="if set multiplies predictions by this number",
            type=float,
            default=1,
        )
        self.parser.add_argument(
            "--ext_disp_to_eval",
            type=str,
            help="optional path to a .npy disparities file to evaluate",
        )
        self.parser.add_argument(
            "--eval_split",
            type=str,
            default="eigen",
            choices=[
                "eigen",
                "odom_9",
                "odom_10",
            ],
            help="which split to run eval on",
        )
        self.parser.add_argument(
            "--save_pred_disps",
            help="if set saves predicted disparities",
            action="store_true",
        )
        self.parser.add_argument(
            "--no_eval", help="if set disables evaluation", action="store_true"
        )
        self.parser.add_argument(
            "--eval_eigen_to_benchmark",
            help="if set assume we are loading eigen results from npy but we want to evaluate using the new benchmark.",
            action="store_true",
        )
        self.parser.add_argument(
            "--eval_out_dir",
            help="if set will output the disparities to this folder",
            type=str,
        )
        self.parser.add_argument(
            "--post_process",
            help="if set will perform the flipping post processing "
            "from the original monodepth paper",
            action="store_true",
        )
        self.parser.add_argument(
            "--zero_cost_volume",
            action="store_true",
            help="If set, during evaluation all poses will be set to 0, and "
            "so we will evaluate the model in single frame mode",
        )
        self.parser.add_argument(
            "--static_camera",
            action="store_true",
            help="If set, during evaluation the current frame will also be"
            "used as the lookup frame, to simulate a static camera",
        )
        self.parser.add_argument(
            "--eval_teacher",
            action="store_true",
            help="If set, the teacher network will be evaluated",
        )

    def parse(self):
        self.options = self.parser.parse_args()
        return self.options
